#!/usr/bin/env python
# -*- coding: utf-8 -*-

# # 1. 从您的主应用和模型文件中导入必要的对象
# from app import app
# from models import db, Supervisor

# print("--- 开始执行 Supervisor 数据查询脚本 ---")

# # 2. 手动创建应用上下文，以便访问数据库
# with app.app_context():
    
#     print("已进入应用上下文，准备查询数据库...")
    
#     # 3. 查询 Supervisor 表中的所有记录
#     all_supervisors = Supervisor.query.all()

#     # 4. 检查是否找到了记录，并打印结果
#     if all_supervisors:
#         count = len(all_supervisors)
#         print(f"\n成功查询到 {count} 条 Supervisor 记录：")
#         print("--------------------------------------------------")
        
#         # 遍历并打印每一条记录的详细信息
#         for index, supervisor in enumerate(all_supervisors, 1):
#             print(f"  记录 #{index}:")
#             print(f"    ID: {supervisor.id}")
#             print(f"    姓名 (Name): {supervisor.name}")
#             print(f"    用户名 (Username): {supervisor.username}")
#             print(f"    邮箱 (Email): {supervisor.email}")
#             print(f"    院系 (Department): {supervisor.department}")
#             print(f"    创建时间 (Created At): {supervisor.created_at.strftime('%Y-%m-%d %H:%M:%S')}")
#             # 打印密码哈希的前15个字符以供参考，而不是完整哈希
#             print(f"    密码哈希 (Password Hash): {supervisor.password_hash[:15]}...")
#             print("--------------------------------------------------")

#     else:
#         print("\n未在 Supervisor 表中找到任何记录。")

# print("--- 脚本执行完毕 ---")


# ldap_user_group_checker.py
import argparse
import sys
from ldap3 import Server, Connection, ALL, MODIFY_REPLACE, MODIFY_DELETE, MODIFY_ADD
from sqlalchemy import create_engine, text

# --- 1. 配置信息 (请根据您的环境修改) ---

# LDAP 服务器配置 (根据您提供的环境信息)
LDAP_SERVER_URI = 'ldap://127.0.0.1:389'
LDAP_BIND_DN = "cn=admin,dc=ai,dc=ustc,dc=edu,dc=cn"
LDAP_BIND_PASSWORD = "Ldap@UstCAi"
LDAP_BASE_DN = "dc=ai,dc=ustc,dc=edu,dc=cn"
LDAP_USER_SEARCH_BASE = f"ou=People,{LDAP_BASE_DN}"
LDAP_GROUP_SEARCH_BASE = f"ou=Group,{LDAP_BASE_DN}"

DATABASE_URI = 'sqlite:////home/user/website/wed-feature-mate/instance/app.db' 

# --- 2. 核心逻辑 ---

def get_supervisor_mapping_from_db(engine):
    """从应用程序数据库获取 '学生用户名 -> 导师名' 的映射."""
    print("Step 1: 从数据库加载学生与导师的正确关系...")
    # status=1 代表 ApplicationStatus.APPROVED
    # user_type=0 代表 UserType.STUDENT
    query = """
        SELECT username, supervisor_name 
        FROM application 
        WHERE status = 1 AND user_type = 0;
    """
    try:
        with engine.connect() as connection:
            result = connection.execute(text(query))
            mapping = {row[0]: row[1] for row in result}
            print(f" -> 成功加载 {len(mapping)} 条已批准的学生记录。")
            return mapping
    except Exception as e:
        print(f"[错误] 无法连接或查询数据库: {e}")
        sys.exit(1)

def get_ldap_data(conn):
    """从LDAP获取所有用户和组的信息."""
    print("Step 2: 从LDAP加载用户、组和成员信息...")
    
    # 获取组信息
    conn.search(LDAP_GROUP_SEARCH_BASE, '(objectClass=posixGroup)', attributes=['cn', 'gidNumber', 'memberUid'])
    groups = conn.entries
    
    group_name_to_gid = {g.cn.value: int(g.gidNumber.value) for g in groups}
    gid_to_group_name = {int(g.gidNumber.value): g.cn.value for g in groups}
    group_members = {g.cn.value: set(g.memberUid.values) if g.memberUid else set() for g in groups}

    # 获取用户信息
    conn.search(LDAP_USER_SEARCH_BASE, '(objectClass=posixAccount)', attributes=['uid', 'gidNumber'])
    users = {u.uid.value: {'gidNumber': int(u.gidNumber.value), 'dn': u.entry_dn} for u in conn.entries}

    print(f" -> 成功加载 {len(users)} 个用户和 {len(groups)} 个组。")
    return users, group_name_to_gid, gid_to_group_name, group_members


def main(dry_run):
    """主执行函数."""
    if dry_run:
        print("\n*** 正在以【演习模式】运行，不会对LDAP进行任何实际修改。 ***\n")
    else:
        print("\n*** 警告：正在以【执行模式】运行，将直接修改LDAP服务器数据！ ***\n")

    # 连接数据库
    db_engine = create_engine(DATABASE_URI)
    supervisor_map = get_supervisor_mapping_from_db(db_engine)

    # 连接LDAP
    try:
        server = Server(LDAP_SERVER_URI, get_info=ALL)
        conn = Connection(server, user=LDAP_BIND_DN, password=LDAP_BIND_PASSWORD, auto_bind=True)
    except Exception as e:
        print(f"[错误] 无法连接到LDAP服务器: {e}")
        sys.exit(1)
        
    try:
        ldap_users, group_name_to_gid, gid_to_group_name, group_members = get_ldap_data(conn)
        
        print("\nStep 3: 开始检查和校对每个用户...\n" + "="*40)
        
        # 遍历LDAP中的每个用户
        for username, user_info in ldap_users.items():
            print(f"\n--- 正在检查用户: {username} ---")

            if username not in supervisor_map:
                print(f"  [跳过] 用户 '{username}' 在LDAP中存在，但在数据库中不是已批准学生。")
                continue

            # 1. 确定正确信息
            correct_supervisor = supervisor_map[username]
            # 遵循您代码中的组名格式化逻辑
            correct_group_name = f"group_{correct_supervisor}"
            
            if correct_group_name not in group_name_to_gid:
                print(f"  [错误] 导师组 '{correct_group_name}' 不存在于LDAP中！请先确保导师账户已正确创建。")
                continue
                
            correct_gid = group_name_to_gid[correct_group_name]
            current_gid = user_info['gidNumber']
            user_dn = user_info['dn']

            # 2. 检查并纠正主组 (gidNumber)
            if current_gid != correct_gid:
                print(f"  [修正] 主组ID不匹配！")
                print(f"    - 当前 GID: {current_gid} ({gid_to_group_name.get(current_gid, '未知或已删除的组')})")
                print(f"    - 正确 GID: {correct_gid} ({correct_group_name})")
                if not dry_run:
                    conn.modify(user_dn, {'gidNumber': [(MODIFY_REPLACE, [str(correct_gid)])]})
                    if conn.result['result'] == 0:
                        print("    => [成功] 已更新用户主组ID。")
                    else:
                        print(f"    => [失败] 更新主组ID失败: {conn.result}")
            else:
                print(f"  [OK] 主组ID正确: {current_gid} ({correct_group_name})")

            # 3. 检查并纠正组成员关系 (一个用户只在一个组)
            found_in_correct_group = False
            for group_name, members in group_members.items():
                group_dn = f"cn={group_name},{LDAP_GROUP_SEARCH_BASE}"
                
                # 如果用户是这个组的成员
                if username in members:
                    if group_name == correct_group_name:
                        found_in_correct_group = True
                        print(f"  [OK] 用户正确地存在于组 '{group_name}' 中。")
                    else:
                        # 从错误的组中移除
                        print(f"  [修正] 用户属于了错误的组 '{group_name}'，将执行移除操作。")
                        if not dry_run:
                            conn.modify(group_dn, {'memberUid': [(MODIFY_DELETE, [username])]})
                            if conn.result['result'] == 0:
                                print(f"    => [成功] 已从 '{group_name}' 组移除。")
                            else:
                                print(f"    => [失败] 移除失败: {conn.result}")
            
            # 如果用户不在正确的组里，则添加进去
            if not found_in_correct_group:
                print(f"  [修正] 用户未在正确的组 '{correct_group_name}' 中，将执行添加操作。")
                if not dry_run:
                    correct_group_dn = f"cn={correct_group_name},{LDAP_GROUP_SEARCH_BASE}"
                    conn.modify(correct_group_dn, {'memberUid': [(MODIFY_ADD, [username])]})
                    if conn.result['result'] == 0:
                        print(f"    => [成功] 已添加至 '{correct_group_name}' 组。")
                    else:
                        print(f"    => [失败] 添加失败: {conn.result}")

    finally:
        conn.unbind()
        print("\n" + "="*40 + "\n检查完成，已断开LDAP连接。")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="校对LDAP用户的主组(gidNumber)和组成员关系，确保与数据库一致。",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        '--execute',
        action='store_true',
        help="""实际执行LDAP修改操作。
如果没有此参数，脚本将以安全的“演习模式”(Dry Run)运行，只报告将要进行的更改。"""
    )
    args = parser.parse_args()

    main(dry_run=not args.execute)